import torch
import torch.nn as nn
import torch.quantization
from torchvision.models.resnet import resnet34

# 请替换为您自己的模型加载方式
# 这里假设您的 ResNet34 模型已经定义并且可以被正确加载
model = ResNet34(feat_dim=80, embed_dim=256, two_emb_layer=False)
model.eval()

# 指定量化配置
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')

# 准备模型进行量化
torch.quantization.prepare(model, inplace=True)

# 假设 calibrate_data_loader 是用于校准的 DataLoader
# 这需要您根据自己的数据集来实现数据加载器
for images, _ in calibrate_data_loader:
    model(images)

# 将模型转换为量化模型
torch.quantization.convert(model, inplace=True)

# 保存量化模型
torch.save(model.state_dict(), 'resnet34_quantized.pth')

print("Quantization complete and model saved.")